// src/client_tests.rs
// Additional unit tests for client.rs to achieve 95%+ coverage

use crate::client::Config;
use tokio::net::{TcpListener, TcpStream};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

// Mock implementations for testing
#[cfg(test)]
mod mock {
    use crate::transport::Transport;
    use async_trait::async_trait;
    use std::io;
    
    pub struct MockTransport {
        pub read_data: Option<Vec<u8>>,
        pub write_data: Vec<u8>,
        pub is_closed_flag: bool,
    }
    
    impl MockTransport {
        pub fn new() -> Self {
            Self {
                read_data: None,
                write_data: Vec::new(),
                is_closed_flag: false,
            }
        }
        
        pub fn with_read_data(data: Vec<u8>) -> Self {
            Self {
                read_data: Some(data),
                write_data: Vec::new(),
                is_closed_flag: false,
            }
        }
        
        pub fn closed() -> Self {
            Self {
                read_data: None,
                write_data: Vec::new(),
                is_closed_flag: true,
            }
        }
    }
    
    #[async_trait]
    impl Transport for MockTransport {
        async fn read(&mut self) -> anyhow::Result<Option<Vec<u8>>> {
            if self.is_closed_flag {
                Ok(None)
            } else {
                Ok(self.read_data.take())
            }
        }
        
        async fn write(&mut self, data: &[u8]) -> anyhow::Result<()> {
            self.write_data.extend_from_slice(data);
            Ok(())
        }
        
        fn is_closed(&self) -> bool {
            self.is_closed_flag
        }

        async fn shutdown(&mut self) -> anyhow::Result<()> {
            self.is_closed_flag = true;
            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use mock::MockTransport;

    async fn setup_test_stream() -> (TcpStream, TcpStream) {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let client_stream = TcpStream::connect(addr).await.unwrap();
        let (server_stream, _) = listener.accept().await.unwrap();
        (client_stream, server_stream)
    }

    #[tokio::test]
    async fn test_socks5_handshake_ipv6_success() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            // We can't directly call the private function, so we'll simulate the handshake here
            let mut stream = server;
            let mut buf = [0u8; 2];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            let nmethods = buf[1];
            let mut methods = vec![0u8; nmethods as usize];
            stream.read_exact(&mut methods).await.unwrap();

            if !methods.contains(&0x00) {
                stream.write_all(&[0x05, 0xFF]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                return Err(anyhow::anyhow!("No supported authentication method"));
            }

            stream.write_all(&[0x05, 0x00]).await.unwrap();

            let mut buf = [0u8; 4];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            if buf[1] != 0x01 { // CONNECT
                stream.write_all(&[0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                return Err(anyhow::anyhow!("Unsupported SOCKS5 command"));
            }

            let (host, port) = match buf[3] { // ATYP
                0x01 => { // IPv4
                    let mut addr_buf = [0u8; 4];
                    stream.read_exact(&mut addr_buf).await?;
                    let host = std::net::Ipv4Addr::from(addr_buf).to_string();
                    let mut port_buf = [0u8; 2];
                    stream.read_exact(&mut port_buf).await?;
                    let port = u16::from_be_bytes(port_buf);
                    (host, port)
                }
                0x03 => { // Domain name
                    let mut len_buf = [0u8; 1];
                    stream.read_exact(&mut len_buf).await?;
                    let len = len_buf[0] as usize;
                    let mut domain_buf = vec![0u8; len];
                    stream.read_exact(&mut domain_buf).await?;
                    let host = String::from_utf8(domain_buf)?;
                    let mut port_buf = [0u8; 2];
                    stream.read_exact(&mut port_buf).await?;
                    let port = u16::from_be_bytes(port_buf);
                    (host, port)
                }
                0x04 => { // IPv6
                    let mut addr_buf = [0u8; 16];
                    stream.read_exact(&mut addr_buf).await?;
                    let host = std::net::Ipv6Addr::from(addr_buf).to_string();
                    let mut port_buf = [0u8; 2];
                    stream.read_exact(&mut port_buf).await?;
                    let port = u16::from_be_bytes(port_buf);
                    (host, port)
                }
                _ => {
                    stream.write_all(&[0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await?;
                    stream.flush().await?;
                    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                    return Err(anyhow::anyhow!("Unsupported address type"));
                }
            };

            // Send reply
            let reply = [0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
            stream.write_all(&reply).await?;

            Ok((stream, host, port))
        });

        // Client side simulation
        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, [0x05, 0x00]);

        // IPv6 address: ::1 (localhost)
        let ipv6_bytes = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1];
        let request = [0x05, 0x01, 0x00, 0x04] // SOCKS5, CONNECT, RSV, IPv6
            .iter()
            .chain(&ipv6_bytes)
            .chain(&[0x00, 0x50]) // port 80
            .cloned()
            .collect::<Vec<u8>>();
        client.write_all(&request).await.unwrap();
        let mut response = [0u8; 10];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response[0..4], [0x05, 0x00, 0x00, 0x01]);

        // Use timeout to prevent hanging
        let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        assert!(result.is_ok());
        let result_value = result.unwrap();
        assert!(result_value.is_ok());
        if let Ok((_, host, port)) = result_value.unwrap() {
            assert_eq!(host, "::1");
            assert_eq!(port, 80);
        }
    }

    #[tokio::test]
    async fn test_socks5_handshake_invalid_version() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            // Simulate the handshake with invalid version
            let mut stream = server;
            let mut buf = [0u8; 2];
            stream.read_exact(&mut buf).await.unwrap();
            
            if buf[0] != 0x05 {
                Err(anyhow::anyhow!("Invalid SOCKS5 version"))
            } else {
                Ok(())
            }
        });

        // Send invalid SOCKS version
        client.write_all(&[0x04, 0x01, 0x00]).await.unwrap();
        
        // Use timeout to prevent hanging
        let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        assert!(result.is_ok());
        let join_result = result.unwrap();
        assert!(join_result.is_ok());
        let task_result = join_result.unwrap();
        assert!(task_result.is_err());
        assert!(task_result.unwrap_err().to_string().contains("Invalid SOCKS5 version"));
    }

    #[tokio::test]
    async fn test_socks5_handshake_no_supported_auth() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            // Simulate the handshake with no supported auth
            let mut stream = server;
            let mut buf = [0u8; 2];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            let nmethods = buf[1];
            let mut methods = vec![0u8; nmethods as usize];
            stream.read_exact(&mut methods).await.unwrap();

            if !methods.contains(&0x00) {
                stream.write_all(&[0x05, 0xFF]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                Err(anyhow::anyhow!("No supported authentication method"))
            } else {
                Ok(())
            }
        });

        // Send methods that don't include 0x00 (no auth)
        client.write_all(&[0x05, 0x01, 0x02]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, [0x05, 0xFF]); // No acceptable methods

        // Use timeout to prevent hanging
        let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        assert!(result.is_ok());
        let join_result = result.unwrap();
        assert!(join_result.is_ok());
        let task_result = join_result.unwrap();
        assert!(task_result.is_err());
        assert!(task_result.unwrap_err().to_string().contains("No supported authentication method"));
    }

    #[tokio::test]
    async fn test_socks5_handshake_unsupported_command() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            // Simulate the handshake with unsupported command
            let mut stream = server;
            let mut buf = [0u8; 2];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            let nmethods = buf[1];
            let mut methods = vec![0u8; nmethods as usize];
            stream.read_exact(&mut methods).await.unwrap();

            if !methods.contains(&0x00) {
                stream.write_all(&[0x05, 0xFF]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                return Err(anyhow::anyhow!("No supported authentication method"));
            }

            stream.write_all(&[0x05, 0x00]).await.unwrap();

            let mut buf = [0u8; 4];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            if buf[1] != 0x01 { // CONNECT
                stream.write_all(&[0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                Err(anyhow::anyhow!("Unsupported SOCKS5 command"))
            } else {
                Ok(())
            }
        });

        // Correct handshake
        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, [0x05, 0x00]);

        // Send unsupported command (0x02 = BIND instead of 0x01 = CONNECT)
        let request = [0x05, 0x02, 0x00, 0x01, 127, 0, 0, 1, 0x00, 0x50];
        client.write_all(&request).await.unwrap();
        let mut response = [0u8; 10];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response[1], 0x07); // Command not supported

        // Use timeout to prevent hanging
        let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        assert!(result.is_ok());
        let join_result = result.unwrap();
        assert!(join_result.is_ok());
        let task_result = join_result.unwrap();
        assert!(task_result.is_err());
        assert!(task_result.unwrap_err().to_string().contains("Unsupported SOCKS5 command"));
    }

    #[tokio::test]
    async fn test_socks5_handshake_unsupported_address_type() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            // Simulate the handshake with unsupported address type
            let mut stream = server;
            let mut buf = [0u8; 2];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            let nmethods = buf[1];
            let mut methods = vec![0u8; nmethods as usize];
            stream.read_exact(&mut methods).await.unwrap();

            if !methods.contains(&0x00) {
                stream.write_all(&[0x05, 0xFF]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                return Err(anyhow::anyhow!("No supported authentication method"));
            }

            stream.write_all(&[0x05, 0x00]).await.unwrap();

            let mut buf = [0u8; 4];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            if buf[1] != 0x01 { // CONNECT
                stream.write_all(&[0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                return Err(anyhow::anyhow!("Unsupported SOCKS5 command"));
            }

            let result = match buf[3] { // ATYP
                0x01 | 0x03 | 0x04 => Ok(()),
                _ => {
                    stream.write_all(&[0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await?;
                    stream.flush().await?;
                    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                    Err(anyhow::anyhow!("Unsupported address type"))
                }
            };

            result
        });

        // Correct handshake
        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, [0x05, 0x00]);

        // Send unsupported address type (0x05)
        let request = [0x05, 0x01, 0x00, 0x05, 127, 0, 0, 1, 0x00, 0x50];
        client.write_all(&request).await.unwrap();
        let mut response = [0u8; 10];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response[1], 0x08); // Address type not supported

        // Use timeout to prevent hanging
        let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        assert!(result.is_ok());
        let join_result = result.unwrap();
        assert!(join_result.is_ok());
        let task_result = join_result.unwrap();
        assert!(task_result.is_err());
        assert!(task_result.unwrap_err().to_string().contains("Unsupported address type"));
    }

    #[tokio::test]
    async fn test_socks5_handshake_domain_name_max_length() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            // Simulate the full handshake with max length domain
            let mut stream = server;
            let mut buf = [0u8; 2];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            let nmethods = buf[1];
            let mut methods = vec![0u8; nmethods as usize];
            stream.read_exact(&mut methods).await.unwrap();

            if !methods.contains(&0x00) {
                stream.write_all(&[0x05, 0xFF]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                return Err(anyhow::anyhow!("No supported authentication method"));
            }

            stream.write_all(&[0x05, 0x00]).await.unwrap();

            let mut buf = [0u8; 4];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            if buf[1] != 0x01 { // CONNECT
                stream.write_all(&[0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                return Err(anyhow::anyhow!("Unsupported SOCKS5 command"));
            }

            let (host, port) = match buf[3] { // ATYP
                0x01 => { // IPv4
                    let mut addr_buf = [0u8; 4];
                    stream.read_exact(&mut addr_buf).await?;
                    let host = std::net::Ipv4Addr::from(addr_buf).to_string();
                    let mut port_buf = [0u8; 2];
                    stream.read_exact(&mut port_buf).await?;
                    let port = u16::from_be_bytes(port_buf);
                    (host, port)
                }
                0x03 => { // Domain name
                    let mut len_buf = [0u8; 1];
                    stream.read_exact(&mut len_buf).await?;
                    let len = len_buf[0] as usize;
                    let mut domain_buf = vec![0u8; len];
                    stream.read_exact(&mut domain_buf).await?;
                    let host = String::from_utf8(domain_buf)?;
                    let mut port_buf = [0u8; 2];
                    stream.read_exact(&mut port_buf).await?;
                    let port = u16::from_be_bytes(port_buf);
                    (host, port)
                }
                0x04 => { // IPv6
                    let mut addr_buf = [0u8; 16];
                    stream.read_exact(&mut addr_buf).await?;
                    let host = std::net::Ipv6Addr::from(addr_buf).to_string();
                    let mut port_buf = [0u8; 2];
                    stream.read_exact(&mut port_buf).await?;
                    let port = u16::from_be_bytes(port_buf);
                    (host, port)
                }
                _ => {
                    stream.write_all(&[0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await?;
                    stream.flush().await?;
                    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                    return Err(anyhow::anyhow!("Unsupported address type"));
                }
            };

            // Send reply
            let reply = [0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
            stream.write_all(&reply).await?;

            Ok((stream, host, port))
        });

        // Correct handshake
        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, [0x05, 0x00]);

        // Domain name with maximum length (255 chars)
        let domain = "a".repeat(255).into_bytes();
        let mut request = vec![0x05, 0x01, 0x00, 0x03, domain.len() as u8];
        request.extend_from_slice(&domain);
        request.extend_from_slice(&[0x01, 0xBB]); // port 443
        client.write_all(&request).await.unwrap();
        let mut response = [0u8; 10];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response[0..4], [0x05, 0x00, 0x00, 0x01]);

        // Use timeout to prevent hanging
        let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        assert!(result.is_ok());
        let result_value = result.unwrap();
        assert!(result_value.is_ok());
        if let Ok((_, host, port)) = result_value.unwrap() {
            assert_eq!(host, "a".repeat(255));
            assert_eq!(port, 443);
        }
    }

    #[tokio::test]
    async fn test_socks5_handshake_empty_domain_name() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            // Simulate the full handshake with empty domain
            let mut stream = server;
            let mut buf = [0u8; 2];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            let nmethods = buf[1];
            let mut methods = vec![0u8; nmethods as usize];
            stream.read_exact(&mut methods).await.unwrap();

            if !methods.contains(&0x00) {
                stream.write_all(&[0x05, 0xFF]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                return Err(anyhow::anyhow!("No supported authentication method"));
            }

            stream.write_all(&[0x05, 0x00]).await.unwrap();

            let mut buf = [0u8; 4];
            stream.read_exact(&mut buf).await.unwrap();
            if buf[0] != 0x05 {
                return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
            }
            if buf[1] != 0x01 { // CONNECT
                stream.write_all(&[0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await.unwrap();
                stream.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                return Err(anyhow::anyhow!("Unsupported SOCKS5 command"));
            }

            let (host, port) = match buf[3] { // ATYP
                0x01 => { // IPv4
                    let mut addr_buf = [0u8; 4];
                    stream.read_exact(&mut addr_buf).await?;
                    let host = std::net::Ipv4Addr::from(addr_buf).to_string();
                    let mut port_buf = [0u8; 2];
                    stream.read_exact(&mut port_buf).await?;
                    let port = u16::from_be_bytes(port_buf);
                    (host, port)
                }
                0x03 => { // Domain name
                    let mut len_buf = [0u8; 1];
                    stream.read_exact(&mut len_buf).await?;
                    let len = len_buf[0] as usize;
                    let mut domain_buf = vec![0u8; len];
                    stream.read_exact(&mut domain_buf).await?;
                    let host = String::from_utf8(domain_buf)?;
                    let mut port_buf = [0u8; 2];
                    stream.read_exact(&mut port_buf).await?;
                    let port = u16::from_be_bytes(port_buf);
                    (host, port)
                }
                0x04 => { // IPv6
                    let mut addr_buf = [0u8; 16];
                    stream.read_exact(&mut addr_buf).await?;
                    let host = std::net::Ipv6Addr::from(addr_buf).to_string();
                    let mut port_buf = [0u8; 2];
                    stream.read_exact(&mut port_buf).await?;
                    let port = u16::from_be_bytes(port_buf);
                    (host, port)
                }
                _ => {
                    stream.write_all(&[0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await?;
                    stream.flush().await?;
                    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                    return Err(anyhow::anyhow!("Unsupported address type"));
                }
            };

            // Send reply
            let reply = [0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
            stream.write_all(&reply).await?;

            Ok((stream, host, port))
        });

        // Correct handshake
        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, [0x05, 0x00]);

        // Empty domain name
        let domain = vec![];
        let mut request = vec![0x05, 0x01, 0x00, 0x03, domain.len() as u8];
        request.extend_from_slice(&domain);
        request.extend_from_slice(&[0x01, 0xBB]); // port 443
        client.write_all(&request).await.unwrap();
        let mut response = [0u8; 10];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response[0..4], [0x05, 0x00, 0x00, 0x01]);

        // Use timeout to prevent hanging
        let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        assert!(result.is_ok());
        let result_value = result.unwrap();
        assert!(result_value.is_ok());
        if let Ok((_, host, port)) = result_value.unwrap() {
            assert_eq!(host, "");
            assert_eq!(port, 443);
        }
    }

    #[tokio::test]
    async fn test_socks5_handshake_incomplete_reads() {
        let (mut client, mut server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            // Try to read only part of the handshake
            let mut buf = [0u8; 1];
            match server.read_exact(&mut buf).await {
                Ok(_) => Err(anyhow::anyhow!("Expected error due to incomplete read")),
                Err(_) => Ok(()), // This is expected
            }
        });

        // Send only part of the handshake
        client.write_all(&[0x05]).await.unwrap();
        
        // Wait a bit to ensure the server tries to read
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
        
        // Close the client connection
        drop(client);

        // Use timeout to prevent hanging
        let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), server_handle).await;
        assert!(result.is_ok());
        let result_value = result.unwrap();
        assert!(result_value.is_ok());
    }

    #[tokio::test]
    async fn test_config_struct() {
        let config = Config {
            listen_addr: "127.0.0.1:1080".to_string(),
            server_addr: "server.example.com:443".to_string(),
            transport: "ws".to_string(),
            no_tls: false,
        };
        
        assert_eq!(config.listen_addr, "127.0.0.1:1080");
        assert_eq!(config.server_addr, "server.example.com:443");
        assert_eq!(config.transport, "ws");
        assert_eq!(config.no_tls, false);
    }

    #[tokio::test]
    async fn test_forward_streams_empty_data() {
        use crate::client::protocol::forward_streams;
        use crate::transport::TcpTransport;
        
        let (stream1, mut stream1_peer) = setup_test_stream().await;
        let (stream2, mut stream2_peer) = setup_test_stream().await;

        let handle = tokio::spawn(async move {
            let transport = TcpTransport::new(stream2);
            forward_streams(stream1, Box::new(transport)).await
        });

        // Close the peer streams to unblock the forward_streams function
        stream1_peer.shutdown().await.unwrap();
        stream2_peer.shutdown().await.unwrap();

        let result = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await;
        assert!(result.is_ok());
        assert!(result.unwrap().unwrap().is_ok());
    }
    
    #[tokio::test]
    async fn test_mock_transport_read_write() {
        use mock::MockTransport;
        use crate::transport::Transport;
        
        let mut mock_transport = MockTransport::with_read_data(b"mock data".to_vec());
        
        // Test reading
        let data = Transport::read(&mut mock_transport).await.unwrap();
        assert!(data.is_some());
        assert_eq!(data.unwrap(), b"mock data");
        
        // Test writing
        Transport::write(&mut mock_transport, b"written data").await.unwrap();
        assert_eq!(mock_transport.write_data, b"written data");
    }
    
    #[tokio::test]
    async fn test_mock_transport_closed() {
        use mock::MockTransport;
        use crate::transport::Transport;
        
        let mock_transport = MockTransport::closed();
        assert!(mock_transport.is_closed());
        
        let mut mock_transport_mut = MockTransport::closed();
        let data = Transport::read(&mut mock_transport_mut).await.unwrap();
        assert!(data.is_none());
    }
    
    #[tokio::test]
    async fn test_mock_transport_empty_read() {
        use mock::MockTransport;
        use crate::transport::Transport;
        
        let mut mock_transport = MockTransport::new();
        let data = Transport::read(&mut mock_transport).await.unwrap();
        assert!(data.is_none());
    }
    
    // Test with different parameter combinations
    #[tokio::test]
    async fn test_socks5_handshake_different_parameters() {
        // Test various combinations of IP addresses and ports
        let test_cases = vec![
            ("127.0.0.1", 80u16),
            ("192.168.1.1", 443u16),
            ("10.0.0.1", 22u16),
            ("0.0.0.0", 8080u16),
        ];
        
        for (ip, port) in test_cases {
            let ipv4 = ip.parse::<std::net::Ipv4Addr>().unwrap();
            let request = [0x05, 0x01, 0x00, 0x01] // SOCKS5, CONNECT, RSV, IPv4
                .iter()
                .chain(&ipv4.octets())
                .chain(&port.to_be_bytes())
                .cloned()
                .collect::<Vec<u8>>();
                
            // Verify the request has the correct structure
            assert_eq!(request.len(), 10);
            assert_eq!(request[0], 0x05); // SOCKS5 version
            assert_eq!(request[1], 0x01); // CONNECT command
            assert_eq!(request[2], 0x00); // Reserved
            assert_eq!(request[3], 0x01); // IPv4 ATYP
            
            // Check the IP address bytes
            assert_eq!(&request[4..8], &ipv4.octets());
            
            // Check the port bytes
            assert_eq!(&request[8..10], &port.to_be_bytes());
        }
    }
    
    #[tokio::test]
    async fn test_socks5_handshake_ipv6_parameters() {
        // Test IPv6 addresses with different ports
        let test_cases = vec![
            ("::1", 80u16),
            ("fe80::1", 443u16),
            ("2001:db8::1", 22u16),
        ];
        
        for (ip_str, port) in test_cases {
            let ipv6 = ip_str.parse::<std::net::Ipv6Addr>().unwrap();
            let request = [0x05, 0x01, 0x00, 0x04] // SOCKS5, CONNECT, RSV, IPv6
                .iter()
                .chain(&ipv6.octets())
                .chain(&port.to_be_bytes())
                .cloned()
                .collect::<Vec<u8>>();
                
            // Verify the request has the correct structure
            assert_eq!(request.len(), 22);
            assert_eq!(request[0], 0x05); // SOCKS5 version
            assert_eq!(request[1], 0x01); // CONNECT command
            assert_eq!(request[2], 0x00); // Reserved
            assert_eq!(request[3], 0x04); // IPv6 ATYP
            
            // Check the IP address bytes
            assert_eq!(&request[4..20], &ipv6.octets());
            
            // Check the port bytes
            assert_eq!(&request[20..22], &port.to_be_bytes());
        }
    }
    
    #[tokio::test]
    async fn test_socks5_handshake_domain_parameters() {
        // Test domain names with different ports
        let long_domain = "a".repeat(255);
        let test_cases = vec![
            ("example.com", 80u16),
            ("google.com", 443u16),
            ("github.com", 22u16),
            (&long_domain, 8080u16), // Maximum length domain
        ];
        
        for (domain, port) in test_cases {
            let domain_bytes = domain.as_bytes();
            let mut request = vec![0x05, 0x01, 0x00, 0x03, domain_bytes.len() as u8];
            request.extend_from_slice(domain_bytes);
            request.extend_from_slice(&port.to_be_bytes());
                
            // Verify the request has the correct structure
            assert_eq!(request[0], 0x05); // SOCKS5 version
            assert_eq!(request[1], 0x01); // CONNECT command
            assert_eq!(request[2], 0x00); // Reserved
            assert_eq!(request[3], 0x03); // Domain name ATYP
            assert_eq!(request[4] as usize, domain_bytes.len()); // Domain length
            
            // Check the domain bytes
            assert_eq!(&request[5..5+domain_bytes.len()], domain_bytes);
            
            // Check the port bytes
            assert_eq!(
                &request[5+domain_bytes.len()..7+domain_bytes.len()], 
                &port.to_be_bytes()
            );
        }
    }
}